#! /usr/bin/env python
# -*- coding: utf-8 -*-
import torch
import time
import cv2
import numpy as np
import pyrealsense2 as rs

print("cuda是否可用：",torch.cuda.is_available())
path="/home/agilex/catkin_ws/src/scout_deeplearning/yolov5_pedestrian"
model = torch.hub.load(path,'custom', path+'/yolov5s.pt', source='local',force_reload=True)
#model=torch.load("best.pt")
# 设置模型为推理模式
model.eval()
# result=model("001.jpg")
# result.show()
#调色板
pose_palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],
                                      [153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],
                                      [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],
                                      [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]],
                                     dtype=np.uint8)
palette_count=len(pose_palette)
#参数设置
#640 480
WIDTH = 1280
HEIGHT = 720
# WIDTH = 640
# HEIGHT = 480
FPS = 30

def detection(frame):
    # 将图像转换为RGB格式
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # 使用Yolov5进行目标检测
    results = model(frame_rgb)
    # 获取检测结果的边界框、类别和置信度
    # boxes = results.xyxy[0].numpy()[:, :4]  # 边界框
    # boxes=boxes.astype(int)# 将小数转换为整数
    # scores = results.xyxy[0].numpy()[:, 4]  # 置信度
    # classes = results.xyxy[0].numpy()[:, 5]  # 类别
    #CUDA 设备上的张量转换为 NumPy 数组
    boxes = results.xyxy[0].cpu().numpy()[:, :4]  # 边界框
    boxes=boxes.astype(int)# 将小数转换为整数
    scores = results.xyxy[0].cpu().numpy()[:, 4]  # 置信度
    classes = results.xyxy[0].cpu().numpy()[:, 5]  # 类别
    # 遍历每个检测结果
    for (x, y, w, h), score, cls in zip(boxes, scores, classes):
        # 绘制边界框和类别标签
        #print(x,y,w,h)
        color=tuple(pose_palette[int(cls)%palette_count].tolist())
        cv2.rectangle(frame, (x, y), (w, h), color, 2)
        text=f'{model.names[int(cls)]} {score:.2f}'
        (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)#获取字体的宽度和高度
        box_coords = ((x, y - text_height), (x + text_width, y))
        cv2.rectangle(frame, box_coords[0], box_coords[1], color, cv2.FILLED)#给文字添加背景
        cv2.putText(frame, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                    (255, 255, 255), 2)
    return frame

def main():
    print("Start Initializing  Camera")
# 配置RealSense相机
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.color, WIDTH, HEIGHT, rs.format.bgr8, FPS)
    # cv2.namedWindow('Yolo5sTest', cv2.WINDOW_NORMAL)
    # cv2.resizeWindow('Yolo5sTest', WIDTH, HEIGHT)
    pipeline.start(config)
    try:
        while True:
            start_time=time.time()
            frames = pipeline.wait_for_frames()
            color_frame = frames.get_color_frame()
            #depth_frame=frame.get_color_frame()#深度图
            if not color_frame:
                continue
            # 将帧转换为 OpenCV 图像
            frame_image = np.asanyarray(color_frame.get_data())
            img=detection(frame_image)
            #计算fps
            fps="fps:"+str(int(1/(time.time()-start_time)))
            cv2.putText(img, fps, (20,20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,0, 255), 2)
            cv2.imshow('Yolo5sTest', img)
           # print("fps:",fps)
            # Press esc or 'q' to close the image window
            if cv2.waitKey(1) & 0xFF==ord('q'):
                cv2.destroyAllWindows()
                print('程序退出！')
                break
    
    finally:
        pipeline.stop()

if __name__ == '__main__':
    main()